///////////////////////////////////////////////////////////////////////
// File:        TwainScanner.cpp
// Description: 
// Author:      Philip Kommnick
// Created:     Mon August 31 16:04:48 CEST 2009
//
// This program is open source. For license terms, see the LICENSE file.
//
///////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "Scanner.h"
#include "TwainScanner.h"

HMODULE TwainScanner::dllHandle = NULL;
UINT TwainScanner::dllRefCnt = 0;
DSMENTRYPROC TwainScanner::pDSMFunc = NULL;
TW_IDENTITY *TwainScanner::appIdentity = NULL;

std::set<TW_UINT32> TwainScanner::openSources;

#define Fix32ToFloat ConvertFix32<Gdiplus::REAL>

typedef std::vector<HWND> WindowList;

TwainScanner::TwainScanner(Sink &s) : Scanner(s), inUse(false)
{
    ZeroMemory(&imageSrc, sizeof(TW_IDENTITY));

    if (++dllRefCnt==1) {
        if (!appIdentity) {
            appIdentity = new TW_IDENTITY;
            ZeroMemory(appIdentity, sizeof(TW_IDENTITY));
            appIdentity->Version.MajorNum = 0;
            appIdentity->Version.MinorNum = 1;
            appIdentity->Version.Language = TWLG_USA;
            appIdentity->Version.Country = TWCY_USA;
            strcpy_s (appIdentity->Version.Info, "3.5");
            appIdentity->ProtocolMajor = TWON_PROTOCOLMAJOR;
            appIdentity->ProtocolMinor = TWON_PROTOCOLMINOR;
            appIdentity->SupportedGroups = DG_IMAGE | DG_CONTROL | DF_APP2;
            strcpy_s (appIdentity->Manufacturer, "TwainPunchCorp");
            strcpy_s (appIdentity->ProductFamily, "Utility");
            strcpy_s (appIdentity->ProductName, "TwainScanner");
        }
        if (!(dllHandle = LoadLibrary(L"twaindsm.dll")) ||
                !(pDSMFunc=(DSMENTRYPROC)GetProcAddress(dllHandle, "DSM_Entry")) ||
				(TWRC_SUCCESS != TwainFunc(NULL, DG_CONTROL, DAT_PARENT, MSG_OPENDSM, (TW_MEMREF)&MainWin::hWnd))) {
            dllRefCnt--;
            FreeLibrary(dllHandle);
            throw std::exception("could not open twain");
        }
    }
}

TwainScanner::~TwainScanner(void)
{
	TwainFunc(NULL, DG_CONTROL, DAT_PARENT, MSG_CLOSEDSM, (TW_MEMREF)&MainWin::hWnd);
    if (--dllRefCnt==0) {
        if (appIdentity) {
            delete appIdentity;
            appIdentity = NULL;
        }
        if (pDSMFunc) {
            pDSMFunc = NULL;
        }
        if (dllHandle) {
            FreeLibrary(dllHandle);
            dllHandle = NULL;
        }
    }
}

bool TwainScanner::SelectDriver(TW_IDENTITY &ident)
{
    TW_IDENTITY tmp;
    if (strlen(ident.Manufacturer) + strlen(ident.ProductFamily) + strlen(ident.ProductName) > 0) {
        ZeroMemory(&tmp, sizeof(TW_IDENTITY));
        if ((TWRC_SUCCESS != TwainFunc(NULL, DG_CONTROL, DAT_IDENTITY, MSG_GETFIRST, &tmp)))
            return false;
        do {
            if (!strcmp(ident.Manufacturer, tmp.Manufacturer) && !strcmp(ident.ProductFamily, tmp.ProductFamily) && !strcmp(ident.ProductName, tmp.ProductName)) {
                memcpy(&ident, &tmp, sizeof(TW_IDENTITY));
                return true;
            }
            ZeroMemory(&tmp, sizeof(TW_IDENTITY));
        } while ((TWRC_ENDOFLIST != TwainFunc(NULL, DG_CONTROL, DAT_IDENTITY, MSG_GETNEXT, &tmp)));
    }
	MainWin::Show(true);
    ZeroMemory(&ident, sizeof(TW_IDENTITY));
    while ((TWRC_SUCCESS == TwainFunc(NULL, DG_CONTROL, DAT_IDENTITY, MSG_USERSELECT, &ident))) {
        if (_strnicmp(ident.ProductName, "WIA", 3))
            return true;
		MessageBox(MainWin::hWnd, L"Bridged WIA drivers are not supported. Please select another driver or cancel driver selection to skip this device.", L"TwainPunch Error", MB_ICONERROR | MB_OK);
        ZeroMemory(&ident, sizeof(TW_IDENTITY));
    }
	MainWin::Show(false);
    return false;
}

void TwainScanner::Configure(HKEY key)
{
    DWORD dataSize = 1024;
    CHAR *data = new CHAR[dataSize];
    if (ERROR_SUCCESS==RegQueryValueEx(key, L"TwainName", NULL, NULL, (LPBYTE)data, &dataSize)) {
        data[dataSize]='\0';
        std::stringstream tokenizer(data);
        std::string token;
        std::getline(tokenizer, token, '|');
        strncpy_s(imageSrc.Manufacturer, sizeof(TW_STR32), token.c_str(), sizeof(TW_STR32));
        std::getline(tokenizer, token, '|');
        strncpy_s(imageSrc.ProductFamily, sizeof(TW_STR32), token.c_str(), sizeof(TW_STR32));
        std::getline(tokenizer, token, '\0');
        strncpy_s(imageSrc.ProductName, sizeof(TW_STR32), token.c_str(), sizeof(TW_STR32));
    }
    if (SelectDriver(imageSrc)) {
        _snprintf_s(data, 1024, _TRUNCATE, "%s|%s|%s", imageSrc.Manufacturer, imageSrc.ProductFamily, imageSrc.ProductName);
        RegSetValueEx(key, L"TwainName", NULL, REG_BINARY, reinterpret_cast<LPBYTE>(data), strlen(data));
    }
    delete[] data;
}

bool TwainScanner::IsActive(void)
{
    return inUse;
}

BOOL CALLBACK TwainScanner::TwainFindWindowHandle(HWND hwnd, LPARAM lParam)
{
	WindowList *list = reinterpret_cast<WindowList*>(lParam);
	if (IsWindow(hwnd) && IsWindowVisible(hwnd) && IsWindowEnabled(hwnd) && hwnd!=MainWin::hWnd) {
		list->push_back(hwnd);
		ShowWindow(hwnd, SW_HIDE);
	}
    return TRUE;
}

TW_UINT16 TwainScanner::TwainFunc(TW_IDENTITY *dest, TW_UINT32 DG,TW_UINT16 DAT,TW_UINT16 MSG, TW_MEMREF pData, TW_STATUS *status)
{
    TW_UINT16 ret;
    ret = pDSMFunc(appIdentity, dest, DG, DAT, MSG, pData);
    if (ret != TWRC_SUCCESS && !status) {
        pDSMFunc(appIdentity, dest, DG_CONTROL, DAT_STATUS, MSG_GET, &status);
    }
    return ret;
}

void TwainScanner::TwainClose(void)
{
    TW_USERINTERFACE twUI;
    twUI.ShowUI = FALSE;
	twUI.hParent = (TW_HANDLE)MainWin::hWnd;
    TwainFunc(&imageSrc, DG_CONTROL, DAT_USERINTERFACE, MSG_DISABLEDS, &twUI);
    TwainFunc(NULL, DG_CONTROL, DAT_IDENTITY, MSG_CLOSEDS, &imageSrc);
    openSources.erase(openSources.find(imageSrc.Id));
    inUse = false;
	MainWin::Show(false);
}

void TwainScanner::Scan(void)
{
    if (openSources.end() != openSources.find(imageSrc.Id))
        throw std::exception("device is in use");
    TW_USERINTERFACE twUI;
    twUI.ShowUI = TRUE;
	twUI.hParent = (TW_HANDLE)MainWin::hWnd;

    if (TWRC_SUCCESS != TwainFunc(NULL, DG_CONTROL, DAT_IDENTITY, MSG_OPENDS, &imageSrc))
        throw std::exception("cannot open source");
    inUse = true;
    try {
		MainWin::Show(true);
        openSources.insert(imageSrc.Id);
        if (!TwainSetCapability(CAP_XFERCOUNT, -1, TRUE))
            throw std::exception("cannot set transfer count");
        if ((TWRC_SUCCESS != TwainFunc(&imageSrc, DG_CONTROL, DAT_USERINTERFACE, MSG_ENABLEDS, (TW_MEMREF)&twUI)))
            throw std::exception("cannot enable source user interface");
    } catch (...) {
        TwainClose();
        throw;
    }
}

bool TwainScanner::TwainSetCapability(TW_UINT16 cap,TW_UINT16 value,BOOL sign)
{
    TW_CAPABILITY twCap;
    pTW_ONEVALUE pVal;
    bool ret = false;

    twCap.Cap = cap;
    twCap.ConType = TWON_ONEVALUE;

    twCap.hContainer = GlobalAlloc(GHND, sizeof(TW_ONEVALUE));
    if (twCap.hContainer) {
        pVal = (pTW_ONEVALUE)GlobalLock(twCap.hContainer);
        pVal->ItemType = sign ? TWTY_INT16 : TWTY_UINT16;
        pVal->Item = (TW_UINT32)value;
        GlobalUnlock(twCap.hContainer);
        ret = (TWRC_SUCCESS == TwainFunc(&imageSrc, DG_CONTROL, DAT_CAPABILITY, MSG_SET, (TW_MEMREF)&twCap));
        GlobalFree(twCap.hContainer);
    }
    return ret;
}

// helper for the for_each loop
void inline UnhideWindow(HWND hWnd)
{
	ShowWindow(hWnd, SW_SHOW);
}

void TwainScanner::TwainProcessMessage()
{
//   TW_INT16 result = TWRC_SUCCESS;
    TW_PENDINGXFERS twPend;
    TW_IMAGEINFO info;
    HANDLE hBitmap = NULL;
    TwainFunc(&imageSrc, DG_CONTROL, DAT_PENDINGXFERS, MSG_GET, (TW_MEMREF)&twPend);
    while (twPend.Count>0) {
        if (TWRC_SUCCESS == TwainFunc(&imageSrc, DG_IMAGE, DAT_IMAGEINFO, MSG_GET, &info)) {
            switch (TwainFunc(&imageSrc, DG_IMAGE, DAT_IMAGENATIVEXFER, MSG_GET, &hBitmap)) {
            case TWRC_XFERDONE:
				{
					// first hide the twain window(s)
					WindowList winList;
					EnumThreadWindows(GetCurrentThreadId(), TwainScanner::TwainFindWindowHandle, reinterpret_cast<LPARAM>(&winList));
					// get the image date from the driver
					BITMAPINFO *bmpHeader = (LPBITMAPINFO)GlobalLock(hBitmap);
					BYTE *dataOffset = reinterpret_cast<LPBYTE>(bmpHeader) + sizeof(BITMAPINFOHEADER);
					if (bmpHeader->bmiHeader.biBitCount <= 8)
						dataOffset += sizeof(RGBQUAD) * static_cast<INT_PTR>((std::pow(2.0, bmpHeader->bmiHeader.biBitCount)));
					// transform the data to gdiplus::bitmap and add dpi information
					Gdiplus::Bitmap imageCpy(bmpHeader, dataOffset);
					imageCpy.SetResolution(Fix32ToFloat(info.XResolution), Fix32ToFloat(info.XResolution));
					GlobalUnlock(bmpHeader);
					// process the image
					try {
						dest.Output(imageCpy);
					} catch (std::exception &e) {
						//TaskDialog(hWnd, tpInstance, MAKEINTRESOURCE(IDS_APPNAME), MAKEINTRESOURCE(IDS_APPQUITINST), MAKEINTRESOURCE(IDS_APPQUITCONT), TDCBF_CLOSE_BUTTON, MAKEINTRESOURCE(IDI_APP), NULL);
						MessageBoxA(NULL, e.what(), "Message Error", MB_ICONERROR | MB_OK);
					}
					// last restore the windows
					std::for_each(winList.begin(), winList.end(), UnhideWindow);
				}
				// no break on purpose!
            case TWRC_CANCEL:
				// free up the image data from the twain driver
                GlobalFree(hBitmap);
                break;
            case TWRC_FAILURE:
				{
	                TW_PENDINGXFERS twPend;
		            TwainFunc(&imageSrc, DG_CONTROL, DAT_PENDINGXFERS, MSG_RESET, &twPend);
			    }
            }
        }
        TwainFunc(&imageSrc, DG_CONTROL, DAT_PENDINGXFERS, MSG_ENDXFER, &twPend);
    }
}

bool TwainScanner::PreTranslate(MSG &msg)
{
    if (inUse) {
        TW_EVENT twEvent;
        twEvent.pEvent = (TW_MEMREF)&msg;
        twEvent.TWMessage = MSG_NULL;
        TW_INT16 retCode = TwainFunc(&imageSrc, DG_CONTROL, DAT_EVENT, MSG_PROCESSEVENT, (TW_MEMREF)&twEvent);

        switch (twEvent.TWMessage) {
        case MSG_XFERREADY:
			TwainProcessMessage();
            break;
        case MSG_CLOSEDSREQ:
            TwainClose();
            break;
            //case MSG_CLOSEDSOK:
            //	break;
        }
        return (retCode == TWRC_DSEVENT);
    }
    return false;
}
